import re
from c2rust.actions.naive_translate import NaiveTranslateCode
from c2rust.actions.naive_add_mainfn import NaiveAddMainFunc
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger
from metagpt.actions import UserRequirement

class NaiveRustTranslator(Role):
    name: str = "NaiveRustTranslator"
    profile: str = "NaiveRustTranslator"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([UserRequirement])
        self.set_actions([NaiveTranslateCode])
        
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
    
class NaiveAddMainFunction(Role):
    name: str = "NaiveAddMainFunction"
    profile: str = "NaiveAddMainFunction"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([NaiveTranslateCode])
        self.set_actions([NaiveAddMainFunc])

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg